import mysql.connector


def get_tables(cnx, database):
    # 查询所有的数据表
    cursor = cnx.cursor()
    query = (f'''
        select TABLE_NAME, TABLE_COMMENT from information_schema.TABLES
        where TABLE_SCHEMA = '{database}'
        ''')
    cursor.execute(query)

    tables = []
    for row in cursor:
        tables.append({'table_name': row[0], 'table_comment': row[1]})
    return tables


def get_column(cnx, database, table_name):
    cursor = cnx.cursor()
    query = (f'''
        select COLUMN_NAME, COLUMN_TYPE, COLUMN_KEY, IS_NULLABLE, COLUMN_DEFAULT, COLUMN_COMMENT
        from information_schema.COLUMNS
        where TABLE_NAME = '{table_name}'
        and TABLE_SCHEMA = '{database}'
        ''')
    cursor.execute(query)
    columns = []
    for row in cursor:
        column_key = '' if row[2] is None else row[2]
        columns.append({'column_name': row[0], 'column_type': row[1], 'column_key': column_key, 'is_nullable': row[3],
                        'column_default': row[4], 'column_comment': row[5]})
    columns = sorted(columns, key=lambda c: c.get('column_name'))
    columns = sorted(columns, key=lambda c: -1 if len(c.get('column_key')) else 0)
    return columns


def get_indices(cnx, database, table_name):
    cursor = cnx.cursor()
    query = (f'''
            select INDEX_NAME, COLUMN_NAME, NON_UNIQUE
            from information_schema.STATISTICS
            where TABLE_NAME = '{table_name}'
            and TABLE_SCHEMA = '{database}'
            ''')
    cursor.execute(query)
    indices = []
    for row in cursor:
        indices.append({'index_name': row[0], 'column_name': row[1], 'non_unique': row[2]})
    return indices


def print_file(tabs_cols):
    file = open("数据字典.md", "w", encoding='utf-8')

    for index, tab_col in enumerate(tabs_cols):
        if index != 0:
            file.write('\n')

        # 输出标题
        tab = tab_col.get('tab')
        tab_comment = tab.get("table_comment") if tab.get('table_comment') is not None else ''
        file.write(f'### {index+1}. {tab.get("table_name")} 【{tab_comment}】 \n')
        file.write('--- \n\n')

        # 输出表结构，表格头+表格内容
        file.write(f'#### {index+1}.1. 表结构 \n')
        file.write('\n')
        file.write('| 字段名 | 数据类型 | 键 | 为空 | 默认值 | 字段描述 | \n')
        file.write('| --- | --- | --- | --- | --- | --- | \n')
        cols = tab_col.get('cols')
        for col in cols:
            column_type = col.get('column_type')

            column_default = col.get('column_default')
            if type(column_default) is bytes:
                column_default = column_default.decode('utf-8')

            column_comment = str(col.get("column_comment"))
            if '\r' in column_comment:
                column_comment = column_comment.replace('\r', '<br>').replace('\n', '')
            else:
                column_comment =column_comment.replace('\n', '<br/>').replace('\r', '')

            col_line = f'| {col.get("column_name")} | {column_type} | {col.get("column_key")} ' \
                       f'| {col.get("is_nullable")} | {column_default} | {column_comment} | \n'
            file.write(col_line)
        file.write('\n')

        # 输出索引
        file.write(f'#### {index + 1}.2. 索引 \n')
        file.write('\n')
        file.write('| 索引名称 | 索引类型 | 字段名 | 可重复 | \n')
        file.write('| --- | --- | --- | --- | \n')
        indices = tab_col.get('indices')
        for inx in indices:
            index_name = inx.get("index_name")
            column_name = inx.get("column_name")

            if index_name in 'PRIMARY':
                index_type = '主键索引'
            elif len(column_name.split(',')) > 1:
                index_type = '联合索引'
            else:
                index_type = '普通索引'

            non_unique = '唯一索引' if inx.get("non_unique") == 0 else '可以重复'

            inx_line = f'| {index_name} | {index_type} | {column_name} | {non_unique} | \n'
            file.write(inx_line)

        file.write('\n')

    file.close()


if __name__ == '__main__':
    # 获取连接
    cnx = mysql.connector.connect(user='root', password='12345678', host='localhost',
                                  database='information_schema')
    # 要查询的库的信息
    database = 'test_db'

    # 获取所有表
    tables = get_tables(cnx, database)

    tabs_cols = []
    # 获取字段
    for tab in tables:
        columns = get_column(cnx, database, tab.get('table_name'))
        indices = get_indices(cnx, database, tab.get('table_name'))
        tabs_cols.append({'tab': tab, 'cols': columns, 'indices': indices})

    # 输出文件
    print_file(tabs_cols)
